#include "cybertron/network/SocketUdpServer.hpp"
#include "cybertron/core/UtilCRT.hpp"
#include "cybertron/core/UtilCRC.hpp"
#include <thread>
#include <asio.hpp>
#include "b64/encode.h"
#include "b64/decode.h"
#include "cybertron/network/perf.hpp"

CYBERTRON_BEGIN

class MessageDecoder
{
public:
	MessageDecoder();

	void onMessage(const void* pBuffer, size_t len);

	void storeBase64Message(std::uint32_t crc, const std::string& b64);

	void removeUnused();

	std::list<std::string>& getMessages()
	{
		return mMessages;
	}

private:
	size_t mMessageStartOffset;
	std::string mReceivedData;
	std::list<std::string> mMessages;
};

MessageDecoder::MessageDecoder()
{
	mMessageStartOffset = (size_t)-1;
}

void MessageDecoder::storeBase64Message(std::uint32_t crc, const std::string& b64)
{
	std::uint32_t crcCal = UtilCRC::crc32_encode(b64.c_str(), (size_t)b64.size());
	if (crcCal != crc)
	{
		return;
	}

	std::stringstream b64Stream;
	b64Stream << b64;
	std::stringstream ss;

	base64::decoder D;
	D.decode(b64Stream, ss);

	mMessages.push_back(ss.str());
}

void MessageDecoder::onMessage(const void* pBuffer, size_t len)
{
	for (size_t i = 0; i < len; ++i)
	{
		mReceivedData.push_back(((const char*)pBuffer)[i]);
	}

	for (;;)
	{
		if (mMessageStartOffset == (size_t)-1)
		{
			// find message start
			if (mReceivedData.size() < 4)
			{
				return;
			}
			for (size_t i = 0; i <= mReceivedData.size() - 4; ++i)
			{
				if (mReceivedData[i] == 1 &&
					mReceivedData[i + 1] == 1 &&
					mReceivedData[i + 2] == 1 &&
					mReceivedData[i + 3] == 1)
				{
					mMessageStartOffset = i + 4;
					break;
				}
			}
		}

		if (mMessageStartOffset == (size_t)-1)
		{
			removeUnused();
			return;
		}

		size_t messageEndOffset = (size_t)-1;
		for (size_t i = mMessageStartOffset + 4; i <= mReceivedData.size() - 4; ++i)
		{
			if (mReceivedData[i] == 0 &&
				mReceivedData[i + 1] == 0 &&
				mReceivedData[i + 2] == 0 &&
				mReceivedData[i + 3] == 0)
			{
				messageEndOffset = i;
				break;
			}
		}

		if (messageEndOffset == (size_t)-1)
		{
			removeUnused();
			return;
		}

		if (messageEndOffset - mMessageStartOffset <= 4)
		{
			mReceivedData.erase(mReceivedData.begin(), mReceivedData.begin() + messageEndOffset + 4);
			mMessageStartOffset = (size_t)-1;
			continue;
		}

		std::uint32_t crc = *((std::uint32_t*)(mReceivedData.c_str() + mMessageStartOffset));
		size_t messageBodyStart = mMessageStartOffset + 4;
		mMessageStartOffset = (size_t)-1;
		std::string messageBody = mReceivedData.substr(messageBodyStart, messageEndOffset - messageBodyStart);
		storeBase64Message(crc, messageBody);
		mReceivedData.erase(mReceivedData.begin(), mReceivedData.begin() + messageEndOffset + 4);
	}
}

void MessageDecoder::removeUnused()
{
	if (mMessageStartOffset == (size_t)-1)
	{
		if (mReceivedData.size() > 1024)
		{
			mReceivedData.erase(mReceivedData.begin(), mReceivedData.begin() + 1020);
		}
	}
	else
	{
		if (mMessageStartOffset > 1024)
		{
			mReceivedData.erase(mReceivedData.begin(), mReceivedData.begin() + mMessageStartOffset);
			mMessageStartOffset = 0;
		}
	}
}

class SocketUdpServerImpl
{
public:
	std::shared_ptr<MessageDecoder> mMessageDecoder;
	std::recursive_mutex mCachedMessageMutex;
	std::list<std::string> mCachedMessages;

	std::shared_ptr<asio::ip::udp::socket> mpSocket;
	std::shared_ptr<asio::io_service> mpIoService;
	std::shared_ptr<bool> mThreadRunning;
	std::shared_ptr<std::thread> mpSenderThread;
	int mNumMaximalMessageCache;
};

CybertronCoreAPI SocketUdpServer::SocketUdpServer()
{
	mpImpl = new SocketUdpServerImpl;

	mpImpl->mNumMaximalMessageCache = -1;
}

CybertronCoreAPI SocketUdpServer::~SocketUdpServer()
{
	delete mpImpl;
	mpImpl = nullptr;
}

CybertronCoreAPI bool SocketUdpServer::initialize(
	int port,
	int numMaximalMessageCache)
{
	clear();

	mpImpl->mMessageDecoder = std::make_shared<MessageDecoder>();
	mpImpl->mNumMaximalMessageCache = numMaximalMessageCache;
	mpImpl->mThreadRunning = std::make_shared<bool>(true);
	mpImpl->mpIoService = std::make_shared<asio::io_service>();
	mpImpl->mpSocket = std::make_shared<asio::ip::udp::socket>(
		*(mpImpl->mpIoService),
		asio::ip::udp::endpoint(asio::ip::udp::v4(), port));

	// start background thread
	auto pThreadRunning = mpImpl->mThreadRunning;

	mpImpl->mpSenderThread = UtilCRT::startThread(
		"UdpReceiver",
		[pThreadRunning, numMaximalMessageCache, this]()
	{
		while (*pThreadRunning)
		{
			asio::ip::udp::endpoint senderEndPoint;

			char buff[1024];
			std::size_t bytes_transferred =
				mpImpl->mpSocket->receive_from(asio::buffer(buff, 1024), senderEndPoint);

			IoPerf::inst().incr_recv_packs(1);
			IoPerf::inst().incr_recv_bytes(bytes_transferred);

			mpImpl->mMessageDecoder->onMessage(buff, bytes_transferred);

			auto& msgs = mpImpl->mMessageDecoder->getMessages();
			if (msgs.size() > 0)
			{
				std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
				{
					auto it = msgs.begin();
					for (; it != msgs.end(); ++it)
					{
						mpImpl->mCachedMessages.push_back((*it));
					}
					msgs.clear();
				}
				if (numMaximalMessageCache > 0)
				{
					while (mpImpl->mCachedMessages.size() > (size_t)numMaximalMessageCache)
					{
						mpImpl->mCachedMessages.pop_front();
					}
				}
			}
		}
	});

	return true;
}

CybertronCoreAPI void SocketUdpServer::clear()
{
	if (mpImpl->mThreadRunning != nullptr)
	{
		(*mpImpl->mThreadRunning) = false;

		if (mpImpl->mpIoService != nullptr)
		{
			mpImpl->mpIoService->stop();
		}
	}

	if (mpImpl->mpSocket != nullptr)
	{
		mpImpl->mpSocket->close();
		mpImpl->mpSocket = nullptr;
	}

	if (mpImpl->mpSenderThread != nullptr)
	{
		mpImpl->mpSenderThread->join();
	}

	mpImpl->mpSenderThread = nullptr;
	mpImpl->mpSocket = nullptr;
	mpImpl->mThreadRunning = nullptr;
	mpImpl->mpIoService = nullptr;
	mpImpl->mNumMaximalMessageCache = -1;
	mpImpl->mMessageDecoder = nullptr;

	std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
	mpImpl->mCachedMessages.clear();
}

CybertronCoreAPI void SocketUdpServer::takeAwayMessages(
	std::list<std::string>& messages)
{
	messages.clear();
	std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
	mpImpl->mCachedMessages.swap(messages);
}

extern void encode(
	const void* pData,
	size_t len,
	std::string& encoded);

CybertronCoreAPI void SocketUdpServer::send(
	const std::string& ip,
	int port,
	const void* pData,
	size_t len)
{
	if (mpImpl->mpSocket == nullptr)
	{
		return;
	}

	std::string encodedMessage;
	encode(pData, len, encodedMessage);

	auto end_point = asio::ip::udp::endpoint(asio::ip::address_v4::from_string(ip), port);

	mpImpl->mpSocket->async_send_to(
		asio::buffer(encodedMessage.c_str(), encodedMessage.size()),
		end_point,
		[](asio::error_code ec, std::size_t bytes_transferred)
	{
		if (!ec) {
			IoPerf::inst().incr_send_packs(1);
			IoPerf::inst().incr_send_bytes(bytes_transferred);
		}
	});
}


///////////////////////////////////////////////////////////////

class SocketUdpRawDataServerImpl
{
public:
    std::recursive_mutex mCachedMessageMutex;
    std::list<char*> mCachedMessages;
    std::list<int> mCachedMessageLens;

    std::shared_ptr<asio::ip::udp::socket> mpSocket;
    std::shared_ptr<asio::io_service> mpIoService;
    std::shared_ptr<bool> mThreadRunning;
    std::shared_ptr<std::thread> mpSenderThread;
    int mNumMaximalMessageCache;
};

CybertronCoreAPI SocketUdpRawDataServer::SocketUdpRawDataServer()
{
    mpImpl = new SocketUdpRawDataServerImpl;

    mpImpl->mNumMaximalMessageCache = -1;
}

CybertronCoreAPI SocketUdpRawDataServer::~SocketUdpRawDataServer()
{
    delete mpImpl;
    mpImpl = nullptr;
}

CybertronCoreAPI bool SocketUdpRawDataServer::initialize(
    int port,
    int numMaximalMessageCache)
{
    clear();

    mpImpl->mNumMaximalMessageCache = numMaximalMessageCache;
    mpImpl->mThreadRunning = std::make_shared<bool>(true);
    mpImpl->mpIoService = std::make_shared<asio::io_service>();
    mpImpl->mpSocket = std::make_shared<asio::ip::udp::socket>(
        *(mpImpl->mpIoService),
        asio::ip::udp::endpoint(asio::ip::udp::v4(), port));

    // start background thread
    auto pThreadRunning = mpImpl->mThreadRunning;

    mpImpl->mpSenderThread = UtilCRT::startThread(
        "UdpReceiver",
        [pThreadRunning, numMaximalMessageCache, this]()
    {
        while (*pThreadRunning)
        {
            asio::ip::udp::endpoint senderEndPoint;

            char buff[1024];
            std::size_t bytes_transferred =
                mpImpl->mpSocket->receive_from(asio::buffer(buff, 1024), senderEndPoint);

            IoPerf::inst().incr_recv_packs(1);
            IoPerf::inst().incr_recv_bytes(bytes_transferred);

            if (bytes_transferred > 0)
            {
                std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
                char* data = new char[bytes_transferred];
                memcpy(data, buff, bytes_transferred);
                mpImpl->mCachedMessages.push_back(data);
                mpImpl->mCachedMessageLens.push_back((int)bytes_transferred);

                if (numMaximalMessageCache > 0)
                {
                    while (mpImpl->mCachedMessages.size() > (size_t)numMaximalMessageCache)
                    {
                        char* data = mpImpl->mCachedMessages.front();
                        delete[] data;
                        mpImpl->mCachedMessages.pop_front();
                        mpImpl->mCachedMessageLens.pop_front();
                    }
                }
            }
        }
    });

    return true;
}

CybertronCoreAPI void SocketUdpRawDataServer::clear()
{
    if (mpImpl->mThreadRunning != nullptr)
    {
        (*mpImpl->mThreadRunning) = false;

        if (mpImpl->mpIoService != nullptr)
        {
            mpImpl->mpIoService->stop();
        }
    }

    if (mpImpl->mpSocket != nullptr)
    {
        mpImpl->mpSocket->close();
        mpImpl->mpSocket = nullptr;
    }

    if (mpImpl->mpSenderThread != nullptr)
    {
        mpImpl->mpSenderThread->join();
    }

    mpImpl->mpSenderThread = nullptr;
    mpImpl->mpSocket = nullptr;
    mpImpl->mThreadRunning = nullptr;
    mpImpl->mpIoService = nullptr;
    mpImpl->mNumMaximalMessageCache = -1;

    std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
    mpImpl->mCachedMessages.clear();
    mpImpl->mCachedMessageLens.clear();
}

CybertronCoreAPI void SocketUdpRawDataServer::clearBuffer()
{
    std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
    mpImpl->mCachedMessages.clear();
    mpImpl->mCachedMessageLens.clear();
}

CybertronCoreAPI void SocketUdpRawDataServer::takeAwayMessage(
    char*& pData, int& len)
{
    std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
    if (!mpImpl->mCachedMessages.empty())
    {
        pData = mpImpl->mCachedMessages.front();
        len = mpImpl->mCachedMessageLens.front();
        mpImpl->mCachedMessages.pop_front();
        mpImpl->mCachedMessageLens.pop_front();
    }
    else
    {
        pData = nullptr;
        len = 0;
    }
}

CybertronCoreAPI void SocketUdpRawDataServer::takeAwayTopMessageAndClear(
    char*& pData, int& len)
{
    std::unique_lock<std::recursive_mutex> lock(mpImpl->mCachedMessageMutex);
    if (!mpImpl->mCachedMessages.empty())
    {
        pData = mpImpl->mCachedMessages.back();
        len = mpImpl->mCachedMessageLens.back();

        mpImpl->mCachedMessages.clear();
        mpImpl->mCachedMessageLens.clear();
    }
    else
    {
        pData = nullptr;
        len = 0;
    }
}

CybertronCoreAPI void SocketUdpRawDataServer::send(
    const std::string& ip,
    int port,
    const void* pData,
    size_t len)
{
    if (mpImpl->mpSocket == nullptr)
    {
        return;
    }

    auto end_point = asio::ip::udp::endpoint(asio::ip::address_v4::from_string(ip), port);

    mpImpl->mpSocket->async_send_to(
        asio::buffer(pData, len),
        end_point,
        [](asio::error_code ec, std::size_t bytes_transferred)
    {
        if (!ec) {
            IoPerf::inst().incr_send_packs(1);
            IoPerf::inst().incr_send_bytes(bytes_transferred);
        }
    });
}
CYBERTRON_END
